package com;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 *
 * Java中如何使用
 * 此处我们主要说一下通过动态生成字节码的方式，其他方式可以自行查找资料。
 *
 * 操作java字节码的工具有两个比较流行，一个是ASM，一个是Javassit 。
 *
 * ASM ：直接操作字节码指令，执行效率高，要是使用者掌握Java类字节码文件格式及指令，对使用者的要求比较高。
 *
 * Javassit 提供了更高级的API，执行效率相对较差，但无需掌握字节码指令的知识，对使用者要求较低。
 *
 * 应用层面来讲一般使用建议优先选择Javassit，如果后续发现Javassit 成为了整个应用的效率瓶颈的话可以再考虑ASM.当然如果开发的是一个基础类库，或者基础平台，还是直接使用ASM吧，相信从事这方面工作的开发者能力应该比较高。
 *
 * Javassit使用方法
 *
 * Javassist是一个开源的分析、编辑和创建Java字节码的类库。是由东京工业大学的数学和计算机科学系的 Shigeru Chiba （千叶 滋）所创建的。它已加入了开放源代码JBoss 应用服务器项目,通过使用Javassist对字节码操作为JBoss实现动态AOP框架。javassist是jboss的一个子项目，其主要的优点，在于简单，而且快速。直接使用java编码的形式，而不需要了解虚拟机指令，就能动态改变类的结构，或者动态生成类。
 *
 * Javassist中最为重要的是ClassPool，CtClass ，CtMethod 以及 CtField这几个类。
 *
 * ClassPool：一个基于HashMap实现的CtClass对象容器，其中键是类名称，值是表示该类的CtClass对象。默认的ClassPool使用与底层JVM相同的类路径，因此在某些情况下，可能需要向ClassPool添加类路径或类字节。
 *
 * CtClass：表示一个类，这些CtClass对象可以从ClassPool获得。
 *
 * CtMethods：表示类中的方法。
 *
 * CtFields ：表示类中的字段。
 *
 * @author wangzhilong
 * @date 2020/4/14
 */
public class MyTransformer implements ClassFileTransformer {

    final static String prefix = "\nlong startTime = System.currentTimeMillis();\n";
    final static String postfix = "\nlong endTime = System.currentTimeMillis();\n";

    //被处理的方法列表
    final static Map<String, List<String>> methodMap = new HashMap<String, List<String>>();


    public MyTransformer() {
        add("com.service.TimeTest.sayHello");
        add("com.service.TimeTest.sayHello2");
    }

    private void add(String methodString) {
        String className = methodString.substring(0, methodString.lastIndexOf("."));
        String methodName = methodString.substring(methodString.lastIndexOf(".") + 1);
        List<String> list = methodMap.get(className);
        if (list == null) {
            list = new ArrayList<String>();
            methodMap.put(className, list);
        }
        list.add(methodName);
    }


    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace("/", ".");
        // 判断加载的class的包路径是不是需要监控的类
        if (methodMap.containsKey(className)) {
            CtClass ctclass = null;
            try {
                //使用全称,用于取得字节码类<使用javassist>
                ctclass = ClassPool.getDefault().get(className);
                for (String methodName : methodMap.get(className)) {
                    String outputStr = "\nSystem.out.println(\"this method " + methodName
                            + " cost:\" +(endTime - startTime) +\"ms.\");";
                    //得到这个方法实例
                    CtMethod ctMethod = ctclass.getDeclaredMethod(methodName);
                    //新定义一个方法叫做比如 sayHello$old
                    String newMethodName = methodName + "$old";
                    //将原来的方法名字修改
                    ctMethod.setName(newMethodName);

                    //构建新的方法，复制原来的方法，名字为原来的名字
                    final CtMethod newMethod = CtNewMethod.copy(ctMethod, methodName, ctclass, null);

                    StringBuilder bodyStr = new StringBuilder();
                    bodyStr.append("{");
                    bodyStr.append(prefix);
                    bodyStr.append(newMethodName + "($$);\n");// 调用原有代码，类似于method();($$)表示所有的参数
                    bodyStr.append(postfix);
                    bodyStr.append(outputStr);
                    bodyStr.append("}");

                    //替换新方法
                    newMethod.setBody(bodyStr.toString());
                    //增加新方法
                    ctclass.addMethod(newMethod);
                }
                return ctclass.toBytecode();
            } catch (NotFoundException | CannotCompileException | IOException e) {
                e.printStackTrace();
            }
        }
        return null;
    }
}
